# tushare_mirror/loader.py
from datetime import datetime, timedelta
from typing import List, Optional
import pandas as pd
import tushare as ts
from vnpy.trader.object import BarData, HistoryRequest
from vnpy.trader.constant import Interval, Exchange
from vnpy_tushare.tushare_datafeed import TushareDatafeed
from .config import Config
from .storage import DataStorage
from .utils import detect_exchange
from .exceptions import APIConnectionError, StorageError
import os
import pandas as pd

class TushareLoader:
    def __init__(self):
        self.datafeed = TushareDatafeed()
        self.storage = DataStorage()
        self._init_tushare()

    def _init_tushare(self) -> None:
        """初始化Tushare Pro接口"""
        ts.set_token(Config.TUSHARE_TOKEN)
        self.pro = ts.pro_api()
        if not self.datafeed.init():
            raise APIConnectionError("Tushare数据源初始化失败")

    def load_daily(
            self,
            symbol: str,
            start: datetime,
            end: datetime
    ) -> List[BarData]:
        """加载日线数据（自动缓存）"""
        # 优先从存储加载
        if self.storage.has_data(symbol, "daily"):
            return self._load_from_storage(symbol, start, end)

        # 从API加载
        req = HistoryRequest(
            symbol=symbol,
            exchange=detect_exchange(symbol),
            start=start,
            end=end,
            interval=Interval.DAILY
        )
        bars = self.datafeed.query_bar_history(req)
        if bars:
            self.storage.save_bars(bars, symbol)
        return bars

    def _load_from_storage(
            self,
            symbol: str,
            start: datetime,
            end: datetime
    ) -> List[BarData]:
        """从本地存储加载日线数据"""
        csv_path = os.path.join(Config.DATA_PATH, f"{symbol}_daily.csv")
        if not os.path.exists(csv_path):
            return []

        df = pd.read_csv(csv_path)
        df["datetime"] = pd.to_datetime(df["datetime"])
        filtered = df[(df["datetime"] >= start) & (df["datetime"] <= end)]
        return [self._convert_df_row_to_bar(row) for _, row in filtered.iterrows()]

    def _convert_df_row_to_bar(self, row: pd.Series) -> BarData:
        """DataFrame行转BarData对象"""
        return BarData(
            symbol=row["symbol"],
            exchange=Exchange(row["exchange"]),
            datetime=row["datetime"],
            interval=Interval(row["interval"]),
            open_price=row["open_price"],
            high_price=row["high_price"],
            low_price=row["low_price"],
            close_price=row["close_price"],
            volume=row["volume"],
            gateway_name="tushare"
        )

    def update_daily_data(self, symbol: str) -> None:
        """自动更新日线数据"""
        latest_bar = self.storage.get_latest_bar(symbol)
        end = datetime.now()
        start = latest_bar.datetime + timedelta(days=1) if latest_bar else end - timedelta(days=365)

        if start <= end:
            bars = self.load_daily(symbol, start, end)
            print(f"更新日线数据 {symbol} 共 {len(bars)} 条")

    def load_valuation(
            self,
            symbol: str,
            start: datetime,
            end: datetime
    ) -> pd.DataFrame:
        """加载估值数据"""
        exchange_code = "SH" if detect_exchange(symbol) == Exchange.SSE else "SZ"
        ts_code = f"{symbol}.{exchange_code}"

        # 从API获取数据
        df = self.pro.daily_basic(
            ts_code=ts_code,
            start_date=start.strftime("%Y%m%d"),
            end_date=end.strftime("%Y%m%d"),
            fields="ts_code,trade_date,pe,pb,ps,dv_ratio,circ_mv"
        )
        if df.empty:
            return pd.DataFrame()

        # 数据格式化
        df = df.rename(columns={"trade_date": "datetime"})
        df["datetime"] = pd.to_datetime(df["datetime"])
        df["symbol"] = symbol
        self.storage.save_valuation(df, symbol)
        return df

    def update_valuation_data(self, symbol: str) -> None:
        """自动更新估值数据"""
        latest_val = self.storage.get_latest_valuation(symbol)
        end = datetime.now()
        start = latest_val["datetime"] + timedelta(days=1) if latest_val else end - timedelta(days=365)

        if start <= end:
            df = self.load_valuation(symbol, start, end)
            print(f"更新估值数据 {symbol} 共 {len(df)} 条")

    def get_latest_valuation(self, symbol: str) -> Optional[dict]:
        """获取最新估值数据（带异常处理）"""
        try:
            return self.storage.get_latest_valuation(symbol)
        except StorageError as e:
            print(f"获取估值数据失败: {str(e)}")
            return None